import csv
import json
import math
import shutil
from pathlib import Path
from typing import List, Dict, Tuple, NamedTuple, Optional
import numpy as np
from shapely.geometry import Point, Polygon
from shapely.ops import unary_union


class PhotoInfo(NamedTuple):
    """照片信息数据结构"""
    path: Path
    longitude: float
    latitude: float
    altitude: float
    timestamp: Optional[float] = None
    easting: Optional[float] = None  # 添加UTM东坐标
    northing: Optional[float] = None  # 添加UTM北坐标


class ClusterCenter(NamedTuple):
    """聚类中心数据结构"""
    name: str
    longitude: float
    latitude: float
    easting: float  # 添加UTM东坐标
    northing: float  # 添加UTM北坐标
    photos: List[PhotoInfo]


def read_gps_data(input_dir: Path) -> List[PhotoInfo]:
    """
    从CSV或JSON文件中读取GPS数据，包括UTM坐标

    参数:
        input_dir: 包含GPS数据的目录

    返回:
        List[PhotoInfo]: 照片信息列表
    """
    photo_info_list = []

    # 首先尝试读取UTM坐标文件
    utm_csv_file = input_dir / "utm6_coordinates.csv"
    if utm_csv_file.exists():
        print("正在读取UTM坐标数据...")
        with open(utm_csv_file, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                photo_info_list.append(PhotoInfo(
                    path=input_dir / row['文件名'],
                    longitude=float(row['原始经度']),
                    latitude=float(row['原始纬度']),
                    altitude=float(row.get('海拔(m)', 0)),
                    timestamp=row.get('拍摄时间', None),
                    easting=float(row['东坐标']),
                    northing=float(row['北坐标'])
                ))
        return photo_info_list

    # 如果没有UTM坐标文件，尝试读取原始GPS坐标文件
    csv_file = input_dir / "gps_coordinates.csv"
    if csv_file.exists():
        print("正在读取GPS坐标数据...")
        with open(csv_file, 'r', encoding='utf-8') as f:
            reader = csv.DictReader(f)
            for row in reader:
                if '纬度' in row and '经度' in row:
                    photo_info_list.append(PhotoInfo(
                        path=input_dir / row['文件名'],
                        longitude=float(row['经度']),
                        latitude=float(row['纬度']),
                        altitude=float(row.get('海拔(m)', 0)),
                        timestamp=row.get('拍摄时间', None)
                    ))
        return photo_info_list

    # 尝试读取JSON文件
    json_file = input_dir / "gps_coordinates.json"
    if json_file.exists():
        print("正在读取JSON格式的GPS数据...")
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
            for item in data:
                photo_info_list.append(PhotoInfo(
                    path=input_dir / item['file'],
                    longitude=item['longitude'],
                    latitude=item['latitude'],
                    altitude=item.get('altitude', 0),
                    timestamp=item.get('timestamp', None)
                ))
        return photo_info_list

    # 如果没有找到GPS数据文件，尝试直接从图片中提取
    print("未找到GPS数据文件，尝试直接从图片中提取...")
    try:
        from gps_extractor import batch_extract_gps

        results = batch_extract_gps(input_dir, "csv")
        for img_path, coord in results.items():
            if coord:
                photo_info_list.append(PhotoInfo(
                    path=img_path,
                    longitude=coord.longitude,
                    latitude=coord.latitude,
                    altitude=coord.altitude,
                    timestamp=coord.timestamp
                ))
    except ImportError:
        print("无法导入gps_extractor模块")

    return photo_info_list


def cluster_photos(photos: List[PhotoInfo], distance: float) -> List[ClusterCenter]:
    """
    根据距离聚类照片（使用UTM坐标）

    参数:
        photos: 照片信息列表
        distance: 聚类距离（米）

    返回:
        List[ClusterCenter]: 聚类中心列表
    """
    if not photos:
        return []

    # 检查是否有UTM坐标
    has_utm_coords = all(photo.easting is not None and photo.northing is not None for photo in photos)

    if not has_utm_coords:
        print("警告: 没有找到UTM坐标，将使用经纬度坐标进行聚类（精度较低）")
        # 创建点集合（使用经纬度）
        points = [Point(photo.longitude, photo.latitude) for photo in photos]

        # 为每个点创建缓冲区（使用简化的度-米转换）
        buffers = []
        for point in points:
            # 简化的度-米转换
            lat_deg_per_m = 1 / 111320
            lon_deg_per_m = 1 / (111320 * math.cos(math.radians(point.y)))

            buffer_distance_lon = distance * lon_deg_per_m
            buffer_distance_lat = distance * lat_deg_per_m

            buffer = point.buffer(max(buffer_distance_lon, buffer_distance_lat))
            buffers.append(buffer)
    else:
        # 使用UTM坐标进行聚类
        print("使用UTM坐标进行精确聚类...")
        # 创建点集合（使用UTM坐标）
        points = [Point(photo.easting, photo.northing) for photo in photos]

        # 为每个点创建缓冲区（使用米作为单位）
        buffers = [point.buffer(distance) for point in points]

    # 计算缓冲区并集
    union_buffer = unary_union(buffers)

    # 如果并集是单个多边形，则直接使用
    if union_buffer.geom_type == 'Polygon':
        clusters = [union_buffer]
    else:
        # 获取所有分离的多边形
        clusters = list(union_buffer.geoms)

    # 为每个聚类找到对应的照片并计算中心点
    cluster_centers = []
    for i, cluster in enumerate(clusters):
        cluster_photos = []
        cluster_longitudes = []
        cluster_latitudes = []
        cluster_eastings = []
        cluster_northings = []

        # 找到属于这个聚类的照片
        for j, point in enumerate(points):
            if cluster.contains(point):
                cluster_photos.append(photos[j])
                cluster_longitudes.append(photos[j].longitude)
                cluster_latitudes.append(photos[j].latitude)

                if has_utm_coords:
                    cluster_eastings.append(photos[j].easting)
                    cluster_northings.append(photos[j].northing)

        if cluster_photos:
            # 计算聚类中心（平均坐标）
            center_longitude = sum(cluster_longitudes) / len(cluster_longitudes)
            center_latitude = sum(cluster_latitudes) / len(cluster_latitudes)

            if has_utm_coords:
                center_easting = sum(cluster_eastings) / len(cluster_eastings)
                center_northing = sum(cluster_northings) / len(cluster_northings)
            else:
                # 如果没有UTM坐标，使用经纬度的平均值
                center_easting = center_longitude
                center_northing = center_latitude

            cluster_centers.append(ClusterCenter(
                name=f"L{i + 1}",
                longitude=center_longitude,
                latitude=center_latitude,
                easting=center_easting,
                northing=center_northing,
                photos=cluster_photos
            ))

    return cluster_centers


def export_locations_csv(cluster_centers: List[ClusterCenter], output_path: Path):
    """
    导出点位信息到CSV文件

    参数:
        cluster_centers: 聚类中心列表
        output_path: 输出文件路径
    """
    with open(output_path, 'w', newline='', encoding='utf-8') as f:
        writer = csv.writer(f)
        writer.writerow(['name', 'longitude', 'latitude'])

        for center in cluster_centers:
            writer.writerow([center.name, center.longitude, center.latitude])


def copy_photos_to_clusters(cluster_centers: List[ClusterCenter], output_dir: Path):
    """
    将照片复制到对应的聚类文件夹

    参数:
        cluster_centers: 聚类中心列表
        output_dir: 输出目录
    """
    # 确保输出目录存在
    output_dir.mkdir(parents=True, exist_ok=True)

    for center in cluster_centers:
        # 创建聚类文件夹
        cluster_dir = output_dir / center.name
        cluster_dir.mkdir(exist_ok=True)

        # 复制照片到聚类文件夹
        for photo in center.photos:
            if photo.path.exists():
                shutil.copy2(photo.path, cluster_dir / photo.path.name)


def run_clustering(input_dir: Path, output_dir: Path, distance: float = 10.0):
    """
    运行聚类算法的主函数

    参数:
        input_dir: 输入目录
        output_dir: 输出目录
        distance: 聚类距离（米）
    """
    # 读取GPS数据
    print("正在读取GPS数据...")
    photos = read_gps_data(input_dir)
    print(f"成功读取 {len(photos)} 张照片的GPS数据")

    # 聚类照片
    print("正在聚类照片...")
    cluster_centers = cluster_photos(photos, distance)
    print(f"找到 {len(cluster_centers)} 个聚类")

    # 导出点位信息
    print("正在导出点位信息...")
    locations_csv = output_dir / "locations.csv"
    export_locations_csv(cluster_centers, locations_csv)

    # 复制照片到聚类文件夹
    print("正在复制照片到聚类文件夹...")
    copy_photos_to_clusters(cluster_centers, output_dir)

    print("处理完成！")
    print(f"点位信息已保存到: {locations_csv}")
    print(f"聚类照片已保存到: {output_dir}")


if __name__ == "__main__":
    # 在PyCharm中测试时，直接设置参数值
    # 如果您想使用命令行参数，请注释掉下面的代码，并取消注释main()调用

    # 直接设置参数值
    input_dir = Path("test_images")  # 替换为您的输入目录
    output_dir = Path("clustered_output")  # 替换为您的输出目录
    distance = 10.0

    run_clustering(input_dir, output_dir, distance)